class KMeans(object):
def __init__(self): # No need to implement
pass
def pairwise_dist(self, image, centers, K):
"""
Args:
image: C x H x W numpy array
centers: K numpy array
K: number of clusters
Return:
dist: N x M array, where dist[i] is the euclidean distance between X[i] and Y[j]
"""
C, H, W = image.shape
dists = np.zeros((K, C, H, W))
for i in range(K):
dists[i] = abs(image-centers[i])
return dists
def _update_assignment(self, image, centers, K):
"""
Args:
image: C x H x W numpy array, the observations
centers: K numpy array, where K is the number of clusters
K: number of clusters
Return:
cluster_idx: C x H x W numpy array, the cluster assignment for each point
"""
cluster_idx = np.empty_like(image)
distances = self.pairwise_dist(image, centers, K)
cluster_idx = np.argmin(distances, axis=0)
return cluster_idx
def _update_centers(self, image, cluster_idx, K):
"""
Args:
image: C x H x W numpy array, the observations
cluster_idx: C x H x W numpy array, the cluster assignment for each point
K: number of clusters
Return:
centers: new centers, K numpy array, where K is the number of clusters
"""
new_centers = np.empty([K])
for i in range(K):
indices = np.where(cluster_idx == i)
if indices[0].size ==0:
print('Center', i, 'is empty!')
new_centers[i] = 0
else:
new_centers[i] = np.mean(image[indices])
return new_centers
def _get_loss(self, image, centers, cluster_idx, K):
"""
Args:
image: C x H x W numpy array, the observations
centers: K numpy array, where K is the number of clusters
cluster_idx: C x H x W numpy array, the cluster assignment for each point
K: number of clusters
Return:
loss: a single float number, which is the objective function of KMeans.
"""
dists = self.pairwise_dist(image, centers, K)
loss = 0.0
for i in range(K):
indices = np.where(cluster_idx == i)
if indices[0].size ==0:
print('Loss', i, 'is zero!')
loss = loss + 0
else:
loss = loss + np.sum(dists[i][indices])
return loss
